# model settings
crop_size = (512, 512)
norm_cfg = dict(type='SyncBN', requires_grad=True)
data_preprocessor = dict(
    type='SegDataPreProcessor',
    mean=[0, 0, 0],
    std=[255, 255, 255],
    size=crop_size,
    bgr_to_rgb=True,
    pad_val=0,
    seg_pad_val=0)
model = dict(
    type='DualBranchSegmentor',   # ✅ 改成自定义的 segmentor
    data_preprocessor=data_preprocessor,
    pretrained='open-mmlab://resnet50_v1c',
    backbone=dict(
        type='ResNetV1c',
        depth=50,
        num_stages=4,
        out_indices=(0, 1, 2, 3),
        dilations=(1, 1, 2, 4),
        strides=(1, 2, 1, 1),
        norm_cfg=norm_cfg,
        norm_eval=False,
        style='pytorch',
        contract_dilation=True),
    decode_head=[
        dict(
            type='FCNHeadWithKey',
            in_channels=2048,
            in_index=3,
            channels=256,
            num_convs=1,
            num_classes=2,
            dropout_ratio=0.1,
            align_corners=False,
            ignore_index=255,
            loss_decode=[
                dict(type='CrossEntropyLoss', loss_weight=1.0, class_weight=[0.1, 0.9]),
                dict(type='DiceLoss', loss_weight=1.0)
            ],
            label_key='gt_seg_map_lcx'
        ),
        dict(
            type='FCNHeadWithKey',
            in_channels=2048,
            in_index=3,
            channels=256,
            num_convs=1,
            num_classes=2,
            dropout_ratio=0.1,
            align_corners=False,
            loss_decode=[
                dict(type='CrossEntropyLoss', loss_weight=1.0, class_weight=[0.1, 0.9]),
                dict(type='DiceLoss', loss_weight=1.0)
            ],
            label_key='gt_seg_map_lad'
        )
    ],
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=1024,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=2,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=[
            dict(type='CrossEntropyLoss', loss_name='loss_aux', loss_weight=0.4, class_weight=[0.1, 0.9]),
            dict(type='DiceLoss', loss_weight=0.4)
        ],
    ),
    train_cfg=dict(),
    test_cfg=dict(mode='whole')
)
